/* *****************************************************************
    MESQUITE -- The Mesh Quality Improvement Toolkit

    Copyright 2004 Sandia Corporation and Argonne National
    Laboratory.  Under the terms of Contract DE-AC04-94AL85000
    with Sandia Corporation, the U.S. Government retains certain
    rights in this software.

    This library is free software; you can redistribute it and/or
    modify it under the terms of the GNU Lesser General Public
    License as published by the Free Software Foundation; either
    version 2.1 of the License, or (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public License
    (lgpl.txt) along with this library; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

    diachin2@llnl.gov, djmelan@sandia.gov, mbrewer@sandia.gov,
    pknupp@sandia.gov, tleurent@mcs.anl.gov, tmunson@mcs.anl.gov

  ***************************************************************** */
/*!
  \file   LPtoPTemplate.cpp
  \brief

  This Objective Function is evaluated using an L P norm to the pth power.
  total=(sum (x_i)^pVal)
  \author Michael Brewer
  \author Thomas Leurent
  \date   2002-01-23
*/
#include <cmath>
#include "LPtoPTemplate.hpp"
#include "MsqFreeVertexIndexIterator.hpp"
#include "MsqTimer.hpp"
#include "MsqHessian.hpp"
#include "MsqDebug.hpp"
#include "QualityMetric.hpp"

using namespace MBMesquite;

LPtoPTemplate::LPtoPTemplate( QualityMetric* qualitymetric, short Pinput, MsqError& err )
    : ObjectiveFunctionTemplate( qualitymetric )
{
    pVal = Pinput;
    if( pVal < 1 )
    {
        MSQ_SETERR( err )( "P_VALUE must be greater than 0.", MsqError::INVALID_ARG );
        return;
    }

    dividingByN = false;

    clear();
}

LPtoPTemplate::LPtoPTemplate( short P, QualityMetric* qm )
    : ObjectiveFunctionTemplate( qm ), pVal( P ), dividingByN( false )
{
    clear();
}

void LPtoPTemplate::clear()
{
    mCount     = 0;
    mPowSum    = 0;
    saveCount  = 0;
    savePowSum = 0;
}

// Michael:  need to clean up here
LPtoPTemplate::~LPtoPTemplate() {}

ObjectiveFunction* LPtoPTemplate::clone() const
{
    return new LPtoPTemplate( *this );
}

double LPtoPTemplate::get_value( double power_sum,
                                 size_t count,
                                 EvalType type,
                                 size_t& global_count,
                                 MsqError& /*err*/ )
{
    double result = 0;
    switch( type )
    {
        default:
        case CALCULATE:
            result       = power_sum;
            global_count = count;
            break;

        case ACCUMULATE:
            mPowSum += power_sum;
            mCount += count;
            result       = mPowSum;
            global_count = mCount;
            break;

        case SAVE:
            savePowSum   = power_sum;
            saveCount    = count;
            result       = mPowSum;
            global_count = mCount;
            break;

        case UPDATE:
            mPowSum -= savePowSum;
            mCount -= saveCount;
            savePowSum = power_sum;
            saveCount  = count;
            mPowSum += savePowSum;
            mCount += saveCount;
            result       = mPowSum;
            global_count = mCount;
            break;

        case TEMPORARY:
            result       = mPowSum - savePowSum + power_sum;
            global_count = mCount + count - saveCount;
            break;
    }

    //  if (!global_count)
    //    {
    //      MSQ_SETERR(err)(" global_count is zero, possibly due to an invalid mesh.",
    //      MsqError::INVALID_MESH); return -1;  // result is invalid
    //    }
    if( dividingByN && global_count ) result /= global_count;
    return result;
}

bool LPtoPTemplate::evaluate( EvalType type, PatchData& pd, double& value_out, bool free, MsqError& err )
{
    QualityMetric* qm = get_quality_metric();
    if( type == ObjectiveFunction::ACCUMULATE )
        qm->get_single_pass( pd, qmHandles, free, err );
    else
        qm->get_evaluations( pd, qmHandles, free, err );
    MSQ_ERRFALSE( err );

    // calculate OF value for just the patch
    std::vector< size_t >::const_iterator i;
    double value, working_sum = 0.0;
    for( i = qmHandles.begin(); i != qmHandles.end(); ++i )
    {
        bool result = qm->evaluate( pd, *i, value, err );
        if( MSQ_CHKERR( err ) || !result ) return false;

        double tmp_val = value;
        for( short j = 1; j < pVal; ++j )
            tmp_val *= value;
        working_sum += fabs( tmp_val );
    }

    // get overall OF value, update member data, etc.
    size_t global_count;
    value_out = qm->get_negate_flag() * get_value( working_sum, qmHandles.size(), type, global_count, err );
    //  if (!global_count)
    //    return false;  // invalid mesh
    //  else
    return true;
}

bool LPtoPTemplate::evaluate_with_gradient( EvalType type,
                                            PatchData& pd,
                                            double& OF_val,
                                            std::vector< Vector3D >& grad_out,
                                            MsqError& err )
{
    QualityMetric* qm = get_quality_metric();
    qm->get_evaluations( pd, qmHandles, OF_FREE_EVALS_ONLY, err );
    MSQ_ERRFALSE( err );

    // zero gradient
    grad_out.clear();
    grad_out.resize( pd.num_free_vertices(), Vector3D( 0.0, 0.0, 0.0 ) );
    bool qm_bool = true;
    double QM_val;
    OF_val = 0.;
    int p1;

    // calculate OF value and gradient for just the patch
    std::vector< size_t >::const_iterator i;
    for( i = qmHandles.begin(); i != qmHandles.end(); ++i )
    {
        qm_bool = qm->evaluate_with_gradient( pd, *i, QM_val, mIndices, mGradient, err );
        if( MSQ_CHKERR( err ) || !qm_bool ) return false;

        QM_val        = fabs( QM_val );
        double QM_pow = 1.0;
        double factor = qm->get_negate_flag();
        if( pVal == 1 )
            QM_pow = 1.0;
        else
        {
            QM_pow = QM_val;
            for( p1 = 2; p1 < pVal; ++p1 )
                QM_pow *= QM_val;
            factor *= QM_pow * pVal;
        }

        OF_val += QM_pow * QM_val;
        for( size_t j = 0; j < mIndices.size(); ++j )
        {
            mGradient[j] *= factor;
            grad_out[mIndices[j]] += mGradient[j];
        }
    }

    // get overall OF value, update member data, etc.
    size_t global_count;
    OF_val = qm->get_negate_flag() * get_value( OF_val, qmHandles.size(), type, global_count, err );
    //  if (!global_count)
    //    return false;  // invalid mesh

    if( dividingByN && global_count )
    {
        const double inv_n = 1.0 / global_count;
        std::vector< Vector3D >::iterator g;
        for( g = grad_out.begin(); g != grad_out.end(); ++g )
            *g *= inv_n;
    }

    return true;
}

bool LPtoPTemplate::evaluate_with_Hessian_diagonal( EvalType type,
                                                    PatchData& pd,
                                                    double& OF_val,
                                                    std::vector< Vector3D >& grad,
                                                    std::vector< SymMatrix3D >& hess_diag,
                                                    MsqError& err )
{
    QualityMetric* qm = get_quality_metric();
    qm->get_evaluations( pd, qmHandles, OF_FREE_EVALS_ONLY, err );
    MSQ_ERRFALSE( err );

    // zero gradient and hessian
    grad.clear();
    grad.resize( pd.num_free_vertices(), 0.0 );
    hess_diag.clear();
    hess_diag.resize( pd.num_free_vertices(), 0.0 );

    double QM_val, QM_pow = 1.0;
    double fac1, fac2;
    const double negate_flag = qm->get_negate_flag();
    bool qm_bool;
    size_t i;
    short p;

    // Loops over all elements in the patch.
    OF_val = 0.0;
    std::vector< size_t >::const_iterator k;
    for( k = qmHandles.begin(); k != qmHandles.end(); ++k )
    {
        // Computes \nabla^2 Q(e). Only the free vertices will have non-zero entries.
        qm_bool = qm->evaluate_with_Hessian_diagonal( pd, *k, QM_val, mIndices, mGradient, mDiag, err );
        if( MSQ_CHKERR( err ) || !qm_bool ) return false;
        QM_val = fabs( QM_val );

        // **** Computes Hessian ****
        const size_t nve = mIndices.size();
        if( pVal == 1 )
        {
            QM_pow = 1.0;
            for( i = 0; i < nve; ++i )
            {
                mDiag[i] *= negate_flag;
                hess_diag[mIndices[i]] += mDiag[i];
            }
            fac1 = 1;
        }
        else if( pVal >= 2 )
        {
            // Computes the coefficients:
            QM_pow = 1.0;
            for( p = 0; p < pVal - 2; ++p )
                QM_pow *= QM_val;
            // 1 - computes p(p-1)Q(e)^{p-2}
            fac2 = pVal * ( pVal - 1 ) * QM_pow;
            // 2 - computes  pQ(e)^{p-1}
            QM_pow *= QM_val;
            fac1 = pVal * QM_pow;

            // fac1 *= qm->get_negate_flag();
            // fac2 *= qm->get_negate_flag();

            for( i = 0; i < nve; ++i )
            {
                SymMatrix3D op( mGradient[i] );
                op *= fac2;
                mDiag[i] *= fac1;
                op += mDiag[i];
                op *= negate_flag;
                hess_diag[mIndices[i]] += op;
            }
        }
        else
        {
            MSQ_SETERR( err )( " invalid P value.", MsqError::INVALID_STATE );
            return false;
        }

        // **** Computes Gradient ****

        // For each vertex in the element ...
        for( i = 0; i < nve; ++i )
        {
            // ... computes p*q^{p-1}*grad(q) ...
            mGradient[i] *= fac1 * negate_flag;
            // ... and accumulates it in the objective function gradient.
            // also scale the gradient by the scaling factor
            assert( mIndices[i] < pd.num_free_vertices() );
            grad[mIndices[i]] += mGradient[i];
        }

        // **** computes Objective Function value \sum_{i=1}^{N_e} |q_i|^P ****
        OF_val += QM_pow * QM_val;
    }

    size_t global_count;
    OF_val = negate_flag * get_value( OF_val, qmHandles.size(), type, global_count, err );
    //  if (!global_count)
    //    return false;  // invalid mesh

    if( dividingByN && global_count )
    {
        const double inv_n = 1.0 / global_count;
        for( i = 0; i < pd.num_free_vertices(); ++i )
        {
            grad[i] *= inv_n;
            hess_diag[i] *= inv_n;
        }
    }

    return true;
}

/*\ For each element, each entry to be accumulated in the Hessian for
    this objective function (\f$ \sum_{e \in E} Q(e)^p \f$ where \f$ E \f$
    is the set of all elements in the patch) has the form:
    \f$ pQ(e)^{p-1} \nabla^2 Q(e) + p(p-1)Q(e)^{p-2} \nabla Q(e) [\nabla Q(e)]^T \f$.

    For \f$ p=2 \f$, this simplifies to
    \f$ 2Q(e) \nabla^2 Q(e) + 2 \nabla Q(e) [\nabla Q(e)]^T \f$.

    For \f$ p=1 \f$, this simplifies to \f$ \nabla^2 Q(e) \f$.

    The \f$ p=1 \f$ simplified version is implemented directly
    to speed up computation.

    This function does not support vertex-based metrics.

    \param pd The PatchData object for which the objective function
           hessian is computed.
    \param hessian this object must have been previously initialized.
*/
bool LPtoPTemplate::evaluate_with_Hessian( EvalType type,
                                           PatchData& pd,
                                           double& OF_val,
                                           std::vector< Vector3D >& grad,
                                           MsqHessian& hessian,
                                           MsqError& err )
{
    QualityMetric* qm = get_quality_metric();
    qm->get_evaluations( pd, qmHandles, OF_FREE_EVALS_ONLY, err );
    MSQ_ERRFALSE( err );
    double negate_flag = qm->get_negate_flag();

    // zero gradient and hessian
    grad.clear();
    grad.resize( pd.num_free_vertices(), 0.0 );
    hessian.zero_out();

    double QM_val, QM_pow = 1.0;
    double fac1, fac2;
    Matrix3D elem_outer_product;
    bool qm_bool;
    size_t i, j, n;
    short p;

    // Loops over all elements in the patch.
    OF_val = 0.0;
    std::vector< size_t >::const_iterator k;
    for( k = qmHandles.begin(); k != qmHandles.end(); ++k )
    {
        // Computes \nabla^2 Q(e). Only the free vertices will have non-zero entries.
        qm_bool = qm->evaluate_with_Hessian( pd, *k, QM_val, mIndices, mGradient, mHessian, err );
        if( MSQ_CHKERR( err ) || !qm_bool ) return false;
        QM_val = fabs( QM_val );

        // **** Computes Hessian ****
        const size_t nve = mIndices.size();
        if( pVal == 1 )
        {
            QM_pow = 1.0;
            n      = 0;
            for( i = 0; i < nve; ++i )
            {
                for( j = i; j < nve; ++j )
                {
                    // negate if necessary
                    mHessian[n] *= negate_flag;
                    hessian.add( mIndices[i], mIndices[j], mHessian[n], err );
                    MSQ_ERRFALSE( err );
                    ++n;
                }
            }
            fac1 = 1;
        }
        else if( pVal >= 2 )
        {
            // Computes the coefficients:
            QM_pow = 1.0;
            for( p = 0; p < pVal - 2; ++p )
                QM_pow *= QM_val;
            // 1 - computes p(p-1)Q(e)^{p-2}
            fac2 = pVal * ( pVal - 1 ) * QM_pow;
            // 2 - computes  pQ(e)^{p-1}
            QM_pow *= QM_val;
            fac1 = pVal * QM_pow;

            // fac1 *= qm->get_negate_flag();
            // fac2 *= qm->get_negate_flag();

            n = 0;
            for( i = 0; i < nve; ++i )
            {
                for( j = i; j < nve; ++j )
                {
                    elem_outer_product.outer_product( mGradient[i], mGradient[j] );
                    elem_outer_product *= fac2;
                    mHessian[n] *= fac1;
                    mHessian[n] += elem_outer_product;
                    mHessian[n] *= negate_flag;
                    hessian.add( mIndices[i], mIndices[j], mHessian[n], err );
                    MSQ_ERRFALSE( err );
                    ++n;
                }
            }
        }
        else
        {
            MSQ_SETERR( err )( " invalid P value.", MsqError::INVALID_STATE );
            return false;
        }

        // **** Computes Gradient ****

        // For each vertex in the element ...
        for( i = 0; i < nve; ++i )
        {
            // ... computes p*q^{p-1}*grad(q) ...
            mGradient[i] *= fac1 * negate_flag;
            // ... and accumulates it in the objective function gradient.
            // also scale the gradient by the scaling factor
            assert( mIndices[i] < pd.num_free_vertices() );
            grad[mIndices[i]] += mGradient[i];
        }

        // **** computes Objective Function value \sum_{i=1}^{N_e} |q_i|^P ****
        OF_val += QM_pow * QM_val;
    }

    size_t global_count;
    OF_val = negate_flag * get_value( OF_val, qmHandles.size(), type, global_count, err );
    //  if (!global_count)
    //    return false;  // invalid mesh

    if( dividingByN && global_count )
    {
        const double inv_n = 1.0 / global_count;
        std::vector< Vector3D >::iterator g;
        for( g = grad.begin(); g != grad.end(); ++g )
            *g *= inv_n;
        hessian.scale( inv_n );
    }

    return true;
}
